import os
import sys
import random


type_dict = {
    "NoneType": 0,
    "医学专科": 1,
    "检查科目": 2,
    "疾病": 3,
    "病毒": 4,
    "症状": 5,
    "细菌": 6,
    "药物": 7
}

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("python3 gen_dataset.py out_train out_eval out_test")
        exit(0)

    TRAINING_FILE = r"./ccks_7_1_competition_data/训练集/entity_type.txt"
    TESTING_FILE = r"./ccks_7_1_competition_data/验证集/entity_validation.txt"

    out_train = sys.argv[1]
    out_eval = sys.argv[2]
    out_test = sys.argv[3]

    out_type_vocabs = "type_vocabs.txt"
    with open(out_type_vocabs, 'w') as f:
        for k, v in type_dict.items():
            f.write(k + "\n")

    lines = 0
    with open(TRAINING_FILE, 'r') as f:
        for _ in f:
            lines += 1

    eval_index = random.sample(range(lines), int(lines * 0.2))
    with open(TRAINING_FILE, 'r') as f, \
        open(out_train, 'w') as f_train, \
        open(out_eval, 'w') as f_eval:
        index = 0
        for line in f:
            print("\rprocessing %d lines." % index, end='')
            f_eval.write(line) if index in eval_index else f_train.write(line)
            index += 1

    with open(TESTING_FILE, 'r') as f, \
        open(out_test, 'w') as f_test:
        for line in f:
            f_test.write(line)

    print(" done. ")

